package com.sheepone.util.structure.tree.binary;

import com.sheepone.util.enums.CompareType;
import com.sheepone.util.structure.node.BinarySearchTreeNode;
import com.sheepone.util.structure.node.RedBlackTreeNode;

import java.util.*;

/**
 * @author Miss.杨
 * @description 二叉搜索树
 * <p/>1. key不能重复
 * <p/>2. 节点的key值必须大于等于左子节点的key值，小于等于右子节点的key值
 * <p/>3. 二叉搜索树的中序遍历就是其按照key值的升序排列
 * <p/>4. 一个节点的前驱节点是没有右子树的
 * <p/>5. 一个节点的后继节点是没有左子树的
 * @since 2024/5/8 - 22:21
 */
public class BinarySearchTree<K extends Comparable<? super K>, V> {
    protected BinarySearchTreeNode<K, V> root;

    protected int size;

    protected void swapNode(RedBlackTreeNode<K, V> node1, RedBlackTreeNode<K, V> node2) {
        K key = node1.getKey();
        V value = node1.getValue();
        node1.setKey(node2.getKey());
        node1.setValue(node2.getValue());
        node2.setKey(key);
        node2.setValue(value);
    }

    protected CompareType compare(K a, K b) {
        int com = a.compareTo(b);
        if (com < 0) {
            return CompareType.LESS_THAN;
        } else if (com > 0) {
            return CompareType.GREATER_THAN;
        }
        return CompareType.EQUAL;
    }

    public int size() {
        return size;
    }

    public boolean isEmpty() {
        return size == 0;
    }

    public void clear() {
        root = null;
        size = 0;
    }


    /**
     * 根据key值查找对应的value值
     *
     * @param key key值
     * @return 对应的value值
     */
    public V get(K key) {
        BinarySearchTreeNode<K, V> node = getNode(key);
        return node == null ? null : node.getValue();
    }


    /**
     * 根据key值查找对应的节点
     *
     * @param key key值
     * @return 对应的节点
     */
    protected BinarySearchTreeNode<K, V> getNode(K key) {
        BinarySearchTreeNode<K, V> node = root;
        while (node != null) {
            CompareType type = compare(node.getKey(), key);
            if (type == CompareType.EQUAL) {
                return node;
            } else if (type == CompareType.LESS_THAN) {
                node = (BinarySearchTreeNode<K, V>) node.getLeft();
            } else {
                node = (BinarySearchTreeNode<K, V>) node.getRight();
            }
        }
        return null;
    }

    public void put(K key, V value) {
        BinarySearchTreeNode<K, V> node = root;
        while (node != null) {
            CompareType type = compare(node.getKey(), key);
            if (type == CompareType.EQUAL) {
                node.setValue(value);
                return;
            } else if (type == CompareType.LESS_THAN) {
                BinarySearchTreeNode<K, V> left = (BinarySearchTreeNode<K, V>) node.getLeft();
                if (left != null) {
                    node = left;
                } else {
                    node.setLeft(new BinarySearchTreeNode<>(key, value));
                    size++;
                    return;
                }
            } else {
                BinarySearchTreeNode<K, V> right = (BinarySearchTreeNode<K, V>) node.getRight();
                if (right != null) {
                    node = right;
                } else {
                    node.setRight(new BinarySearchTreeNode<>(key, value));
                    size++;
                    return;
                }
            }
        }
        root = new BinarySearchTreeNode<>(key, value);
        size++;
    }

    /**
     * 获取当前传递过来的root节点及其子节点中的最小key值的节点
     *
     * @param root 获取当前传递过来的root节点
     * @return root节点及其子节点中的最小key值的节点
     */

    protected BinarySearchTreeNode<K, V> min(BinarySearchTreeNode<K, V> root) {
        if (root == null) {
            return null;
        }
        BinarySearchTreeNode<K, V> node = root;
        while (node.getLeft() != null) {
            node = (BinarySearchTreeNode<K, V>) node.getLeft();
        }
        return node;
    }

    /**
     * 获取当前传递过来的root节点及其子节点中的最大key值的节点
     *
     * @param root 获取当前传递过来的root节点
     * @return root节点及其子节点中的最大key值的节点
     */

    protected BinarySearchTreeNode<K, V> max(BinarySearchTreeNode<K, V> root) {
        if (root == null) {
            return null;
        }
        BinarySearchTreeNode<K, V> node = root;
        while (node.getRight() != null) {
            node = (BinarySearchTreeNode<K, V>) node.getRight();
        }
        return node;
    }

    /**
     * 获取当前传递过来的节点的前驱节点
     * 1. 当前节点有左子节点，则前驱节点为左子树中的最大节点
     * 2. 当前节点没有左子节点，则前驱节点为当前节点的最近祖先节点，该最近祖先节点为null或者当前节点为该最近祖先节点的右子树的节点
     *
     * @param key 当前节点的key值
     * @return 当前节点的前驱节点
     */
    protected BinarySearchTreeNode<K, V> predecessor(K key) {
        BinarySearchTreeNode<K, V> node = root;
        BinarySearchTreeNode<K, V> predecessor = null;
        while (node != null) {
            if (node.getKey().compareTo(key) == 0) {
                break;
            } else if (node.getKey().compareTo(key) > 0) {
                node = (BinarySearchTreeNode<K, V>) node.getLeft();
            } else {
                predecessor = node;
                node = (BinarySearchTreeNode<K, V>) node.getRight();
            }
        }
        if (node == null) {
            return null;
        }
        if (node.getLeft() != null) {
            return max((BinarySearchTreeNode<K, V>) node.getLeft());
        }
        return predecessor;
    }

    /**
     * 获取当前传递过来的节点的后继节点
     * 1. 当前节点有右子节点，则后继节点为右子树中的最小节点
     * 2. 当前节点没有右子节点，则后继节点为当前节点的最近祖先节点，该最近祖先节点为null或者当前节点为该最近祖先节点的左子树的节点
     *
     * @param key 当前节点的key值
     * @return 当前节点的后继节点
     */
    protected BinarySearchTreeNode<K, V> successor(K key) {
        BinarySearchTreeNode<K, V> node = root;
        BinarySearchTreeNode<K, V> successor = null;
        while (node != null) {
            if (node.getKey().compareTo(key) == 0) {
                break;
            } else if (node.getKey().compareTo(key) > 0) {
                successor = node;
                node = (BinarySearchTreeNode<K, V>) node.getLeft();
            } else {
                node = (BinarySearchTreeNode<K, V>) node.getRight();
            }
        }
        if (node == null) {
            return null;
        }
        if (node.getRight() != null) {
            return min((BinarySearchTreeNode<K, V>) node.getRight());
        }
        return successor;
    }


    /**
     * 查找小于等于key值的第一个节点
     *
     * @param key
     * @return
     */
    protected BinarySearchTreeNode<K, V> findFirstLeNode(K key) {
        BinarySearchTreeNode<K, V> node = root;
        while (node != null) {
            CompareType type = compare(node.getKey(), key);
            // 如果相等或者小于key值，则返回当前节点
            if (type == CompareType.EQUAL || type == CompareType.LESS_THAN) {
                return node;
            } else {
                node = (BinarySearchTreeNode<K, V>) node.getRight();
            }
        }
        return null;
    }

    /**
     * 查找大于等于key值的第一个节点
     *
     * @param key
     * @return
     */
    protected BinarySearchTreeNode<K, V> findFirstGeNode(K key) {
        BinarySearchTreeNode<K, V> node = root;
        while (node != null) {
            CompareType type = compare(node.getKey(), key);
            // 如果相等或者大于key值，则返回当前节点
            if (type == CompareType.EQUAL || type == CompareType.GREATER_THAN) {
                return node;
            } else {
                node = (BinarySearchTreeNode<K, V>) node.getLeft();
            }
        }
        return null;
    }

    /**
     * 利用层序遍历获取小于等于key值的所有节点
     *
     * @param key key值
     * @return 小于等于key值的所有节点
     */
    public List<V> le(K key) {
        BinarySearchTreeNode<K, V> firstLeNode = findFirstLeNode(key);
        if (firstLeNode == null) {
            return Collections.emptyList();
        }
        Queue<BinarySearchTreeNode<K, V>> queue = new LinkedList<>();
        queue.offer(firstLeNode);
        List<V> list = new LinkedList<>();
        while (!queue.isEmpty()) {
            BinarySearchTreeNode<K, V> node = queue.poll();
            list.add(node.getValue());
            if (node.getLeft() != null) {
                queue.offer((BinarySearchTreeNode<K, V>) node.getLeft());
            }
            if (node.getRight() != null) {
                queue.offer((BinarySearchTreeNode<K, V>) node.getRight());
            }
        }
        return list;
    }

    /**
     * 利用层序遍历获取大于等于key值的所有节点
     *
     * @param key key值
     * @return 大于等于key值的所有节点
     */
    public List<V> ge(K key) {
        BinarySearchTreeNode<K, V> firstGeNode = findFirstGeNode(key);
        if (firstGeNode == null) {
            return Collections.emptyList();
        }
        Queue<BinarySearchTreeNode<K, V>> queue = new LinkedList<>();
        queue.offer(firstGeNode);
        List<V> list = new LinkedList<>();
        while (!queue.isEmpty()) {
            BinarySearchTreeNode<K, V> node = queue.poll();
            list.add(node.getValue());
            if (node.getLeft() != null) {
                queue.offer((BinarySearchTreeNode<K, V>) node.getLeft());
            }
            if (node.getRight() != null) {
                queue.offer((BinarySearchTreeNode<K, V>) node.getRight());
            }
        }
        return list;
    }

    /**
     * 利用层序遍历获取小于key值的所有节点
     *
     * @param key key值
     * @return 小于key值的所有节点
     */
    public List<V> lt(K key) {
        BinarySearchTreeNode<K, V> firstLeNode = findFirstLeNode(key);
        if (firstLeNode == null) {
            return Collections.emptyList();
        }
        if (compare(firstLeNode.getKey(), key) == CompareType.EQUAL) {
            firstLeNode = (BinarySearchTreeNode<K, V>) firstLeNode.getLeft();
        }
        if (firstLeNode == null) {
            return Collections.emptyList();
        }
        Queue<BinarySearchTreeNode<K, V>> queue = new LinkedList<>();
        queue.offer(firstLeNode);
        List<V> list = new LinkedList<>();
        while (!queue.isEmpty()) {
            BinarySearchTreeNode<K, V> node = queue.poll();
            list.add(node.getValue());
            if (node.getLeft() != null) {
                queue.offer((BinarySearchTreeNode<K, V>) node.getLeft());
            }
            if (node.getRight() != null) {
                queue.offer((BinarySearchTreeNode<K, V>) node.getRight());
            }
        }
        return list;
    }

    /**
     * 利用层序遍历获取大于key值的所有节点
     *
     * @param key key值
     * @return 大于key值的所有节点
     */
    public List<V> gt(K key) {
        BinarySearchTreeNode<K, V> firstGeNode = findFirstGeNode(key);
        if (firstGeNode == null) {
            return Collections.emptyList();
        }
        if (compare(firstGeNode.getKey(), key) == CompareType.EQUAL) {
            firstGeNode = (BinarySearchTreeNode<K, V>) firstGeNode.getRight();
        }
        if (firstGeNode == null) {
            return Collections.emptyList();
        }
        Queue<BinarySearchTreeNode<K, V>> queue = new LinkedList<>();
        queue.offer(firstGeNode);
        List<V> list = new LinkedList<>();
        while (!queue.isEmpty()) {
            BinarySearchTreeNode<K, V> node = queue.poll();
            list.add(node.getValue());
            if (node.getLeft() != null) {
                queue.offer((BinarySearchTreeNode<K, V>) node.getLeft());
            }
            if (node.getRight() != null) {
                queue.offer((BinarySearchTreeNode<K, V>) node.getRight());
            }
        }
        return list;
    }

    /**
     * 利用中序遍历查找key值在[lk,rk]之间的所有节点
     *
     * @param lk 区间左边界
     * @param rk 区间右边界
     * @return key值在[lk, rk]之间的所有节点
     */
    public List<V> between(K lk, K rk) {
        BinarySearchTreeNode<K, V> node = root;
        Stack<BinarySearchTreeNode<K, V>> stack = new Stack<>();
        List<V> list = new LinkedList<>();
        while (node != null || !stack.isEmpty()) {
            while (node != null) {
                stack.push(node);
                node = (BinarySearchTreeNode<K, V>) node.getLeft();
            }
            if (!stack.isEmpty()) {
                node = stack.pop();
                CompareType rt = compare(node.getKey(), rk);
                if (rt == CompareType.GREATER_THAN) {
                    break;
                }
                CompareType lt = compare(node.getKey(), lk);
                if ((lt == CompareType.GREATER_THAN || lt == CompareType.EQUAL)) {
                    list.add(node.getValue());
                }
                node = (BinarySearchTreeNode<K, V>) node.getRight();
            }
        }
        return list;
    }

    /**
     * 删除指定key的对应节点
     * <p/>1. 删除节点没有左孩子or没有右孩子or左右孩子都没有，将其孩子托孤给parent
     * <p/>2. 删除节点左右孩子都有，用该节点的后继节点替换该节点即可，如果当前节点与其后继节点不直接相邻需要将该后继节点的右孩子托孤给该后继节点的父亲作为其父亲的左孩子
     *
     * @param key 指定key
     */
    public void remove(K key) {
        BinarySearchTreeNode<K, V> node = root;
        BinarySearchTreeNode<K, V> parent = null;
        while (node != null) {
            CompareType type = compare(node.getKey(), key);
            if (type == CompareType.EQUAL) {
                break;
            }
            parent = node;
            if (type == CompareType.LESS_THAN) {
                node = (BinarySearchTreeNode<K, V>) node.getLeft();
            } else {
                node = (BinarySearchTreeNode<K, V>) node.getRight();
            }
        }
        // 没有找到该节点，不需要执行删除操作，直接返回
        if (node == null) {
            return;
        }
        BinarySearchTreeNode<K, V> left = (BinarySearchTreeNode<K, V>) node.getLeft();
        BinarySearchTreeNode<K, V> right = (BinarySearchTreeNode<K, V>) node.getRight();
        // 1.3 删除节点左右孩子都没有，将其孩子托孤给parent
        BinarySearchTreeNode<K, V> child = null;
        // 2 删除节点左右孩子都有，用该节点的后继节点替换该节点即可
        if (left != null && right != null) {
            // 存在右子树，直接在右子树找最小值作为当前节点的后继即可
            child = right;
            BinarySearchTreeNode<K, V> sp = node;
            while (child.getLeft() != null) {
                sp = child;
                child = (BinarySearchTreeNode<K, V>) child.getLeft();
            }
            // 将当前节点的左孩子托孤给后继节点作为其后继节点的左孩子
            child.setLeft(left);
            // 如果当前节点与其后继节点不直接相邻需要将该后继节点的右孩子托孤给该后继节点的父亲作为其父亲的左孩子
            if (right != child) {
                sp.setLeft(child.getRight());
                child.setRight(node.getRight());
            }
        } else {
            if (left != null) {
                // 1.1 删除节点有左孩子，将左孩子托孤给parent
                child = left;
            } else if (right != null) {
                // 1.2 删除节点有右孩子，将右孩子托孤给parent
                child = right;
            }
        }
        // 被删除节点为根节点
        if (parent == null) {
            root = child;
        } else {
            // 被删除节点是左孩子就放左边，被删除节点是右孩子就放右边
            if (parent.getLeft() == node) {
                parent.setLeft(child);
            } else {
                parent.setRight(child);
            }
        }
        size--;
    }
}
